import numpy as np
from python.lib import json
from python.lib.heapq import heappush, heappop

from topo1 import build_hybrid_topology_shared_tor


class StaticRouting:
    def __init__(self, network):
        self.network = network  # 网络拓扑对象
        self.routes = {}  # 预计算的路由表 {(src, dst): path}

    def precompute_routes(self, demand_matrix):
        """
        离线预计算最短时延路径（Dijkstra算法）
        """
        # 遍历所有的流量需求对
        for src in demand_matrix:
            for dst in demand_matrix[src]:
                if src != dst and demand_matrix[src][dst] > 0: # 如果源节点和目标节点不相等且流量大于0
                    # 调用Dijkstra算法计算最短时延路径
                    path = self.dijkstra(src, dst)
                    # 定义路由表
                    self.routes[(src, dst)] = path

    def dijkstra(self, src, dst):
        """
        考虑光链路激活状态的改进Dijkstra算法
        """
        heap = [(0, src, [])] # 优先级队列，存储(cost, node, path)
        visited = set() # 用于记录已访问过的节点


        while heap:
            cost, u, path = heappop(heap)
            if u in visited:
                continue
            visited.add(u)

            if u == dst:
                return path

            for v in self.network.topology[u]:
                # 获取链路当前状态下的时延
                link_delay = self.network.get_link_delay(u, v, 0)  # 静态路由不考虑流量
                new_cost = cost + link_delay
                heappush(heap, (new_cost, v, path + [(u, v)]))

        return None  # 无可用路径

    def calculate_max_delay(self, demand_matrix):
        max_delay = 0
        for (src, dst), path in self.routes.items():
            total_delay = 0
            for u, v in path:
                total_delay += self.network.get_link_delay(u, v, demand_matrix[src][dst])
            max_delay = max(max_delay, total_delay)
        return max_delay

if __name__ == '__main__':
    # 创建测试拓扑
    G = build_hybrid_topology_shared_tor(num_tor=4)  # 从函数中获取拓扑
    print("图的节点G.nodes:", G.nodes())
    print("===========输出图的边===========")
    print("图的边G.edges:", G.edges())
    # 输出每条边的链路容量
    # for u, v, data in G.edges(data=True):
    #     print(f"链路 {u}→{v} 的容量为：{data['capacity']}")

    # 从json文件中读取流量数据
    with open('traffic_simulation.json', 'r') as f:
        data = json.load(f)
        # 输出时间为0.12时候的流量需求矩阵作为all_to_all
        for record in data:
            if record['timestamp'] == 0.12:
                all_to_all_list = record['matrix']
                # 转换为矩阵形式
                all_to_all = np.array(all_to_all_list)
                print(all_to_all)
            else:
                # 输出时间为5.19时候的流量需求矩阵作为brust
                if record['timestamp'] == 5.19:
                    brust_list = record['matrix']
                    brust = np.array(brust_list)
                    print(brust)
    print("输出的all_to_all流量需求矩阵为:", all_to_all)
    print("输出的brust流量需求矩阵为:", brust)
    hybrid_flow = brust
    print("输出的混合流量需求矩阵为:", hybrid_flow)

    # 根据流量需求矩阵获取通信对
    all_to_all_dict = {(i, j): all_to_all[i][j] for i in range(all_to_all.shape[0]) for j in range(all_to_all.shape[1])}
    hybrid_flow_dict = {(i, j): hybrid_flow[i][j] for i in range(hybrid_flow.shape[0]) for j in
                        range(hybrid_flow.shape[1])}
    # 获取all_to_all通信对
    T_all_to_all = []
    for pair, demand in all_to_all_dict.items():
        # print("pair:",pair)
        # print("demand:",demand)
        if demand > 0:
            T_all_to_all.append(pair)
    for pair in T_all_to_all:
        print("pair:", pair)  # 输出的通信对的格式为（1,2），表示ToR1和ToR2

    # 将键从 (i, j) 转换为 ("ToR{i}", "ToR{j}")
    all_to_all_str_keys = {(f"ToR{i + 1}", f"ToR{j + 1}"): value for (i, j), value in all_to_all_dict.items()}
    # print("D_str_keys:",D_str_keys)

    # 将值从 numpy.ndarray 转换为浮点数
    all_to_all_float_values = {k: float(v) for k, v in all_to_all_str_keys.items()}

    # 生成通信对列表
    T_all_to_all = [pair for pair, demand in all_to_all_float_values.items() if demand > 0]
    print("生成all_to_all通信对列表", T_all_to_all)

    # 生成all_to_all+brust通信对
    T_hybird_flow = []
    for pair, demand in hybrid_flow_dict.items():
        # print("pair:",pair)
        # print("demand:",demand)
        if demand > 0:
            T_hybird_flow.append(pair)
    # for pair in T_hybird_flow:
    #     print("pair:", pair)  # 输出的通信对的格式为（1,2），表示ToR1和ToR2

    # 将键从 (i, j) 转换为 ("ToR{i}", "ToR{j}")
    hybrid_str_keys = {(f"ToR{i + 1}", f"ToR{j + 1}"): value for (i, j), value in hybrid_flow_dict.items()}
    # print("D_str_keys:",D_str_keys)

    # 将值从 numpy.ndarray 转换为浮点数
    hybrid_float_values = {k: float(v) for k, v in hybrid_str_keys.items()}

    # 生成通信对列表
    T_hybird_flow = [pair for pair, demand in hybrid_float_values.items() if demand > 0]
    print("生成通信对列表", T_hybird_flow)

    # 计算基础流量路径
    routing = StaticRouting(G)
    precompute_routes = routing.precompute_routes(all_to_all_float_values)
    # print("计算基础流量路径", precompute_routes)
    for src, dst in T_all_to_all:
        print(f"{src} -> {dst}: {precompute_routes[(src, dst)]}")
